#include <CType.h>
#include <stdio.h>
#include <syslib.h>
#include <conio.h>
#include <string.h>
#include <asmlib.h>

IMAGE_DOS_HEADER idh;
IMAGE_FILE_HEADER ifh;
IMAGE_OPTIONAL_HEADER ioh;
IMAGE_SECTION_HEADER pish;

int16_t printf(int8_t * fmt, ...)
{
    int8_t * varhead;
    int8_t p8;
    int16_t p16;
    int32_t p32;
    int8_t * ps;
    int16_t upf = 0, lvf = 0;
    int16_t i = 0, p = 0, l, n = 0;
    int8_t cl, cr, pl;
    int16_t lpv;

    l = strlen(fmt);

    varhead = ((int8_t *)&fmt) + sizeof(int8_t *);

    while(p < l)
        {
        if (fmt[i] != '%')
            {
            switch(fmt[i])
                {
                case '\n':
                case '\r':
                    if ((i < l - 1) && (fmt[i+1] == '\r' || fmt[i+1] == '\n'))
                        {
                        outCStr(CRLF);
                        i += 2;
                        p += 2;
                        n += 2;
                        continue;
                        }
                    else
                        {
                        outCStr(CRLF);
                        n++;
                        }

                    break;
                case '\t':
                    getCursorPos(&cr, &cl);

                    if (cl % 8 == 0)
                        pl = 8;
                    else
                        pl = 8 - (cl % 8);

                    for (lpv = 0; lpv < pl; lpv++)
                        {
                        putCh(' ');
                        }
                    n += pl;
                    i++, p++;
                    continue;
                default:
                    putCh(fmt[i]);
                    break;
                }
            i++, p++, n++;
            continue;
            }
        else if (fmt[i] == '%' && i < l - 1 && fmt[i+1] == '%')
            {
            i += 2;
            p += 2;
            n++;
            putCh('%');
            }
        else
            {
            i++;
            lvf = upf = 0;
            if (fmt[i] == 'l')
                {
                lvf = 1;
                i++;
                }
            switch(fmt[i])
                {
                case 'c':
                    p16 = *(int16_t *)varhead;

                    varhead = (int8_t *)varhead + sizeof(int16_t);
                    p8 = p16 & 0xFF;

                    if (p8 == '\n')
                        putCh('\r');

                    putCh(p8);
                    n++;
                    break;
                case 'd':
                case 'i':
                    if (lvf == 0)
                        {
                        p16 = *(int16_t *)varhead;
                        varhead = (int8_t *)varhead + sizeof(int16_t);
                        n += outCStr(itoa(p16));
                        }
                    else
                        {
                        p32 = *(int32_t *)varhead;
                        varhead = (int8_t *)varhead + sizeof(int32_t);
                        n += outCStr(ltoa(p32));
                        }
                    break;
                case 's':
                    ps = *(int8_t **)varhead;
                    varhead = (int8_t *)(&(((int8_t **)varhead)[1]));
                    n += outCStr(ps);
                    break;
                case 'p':
                    outCStr("0x");
                case 'X':
                    upf = 1;
                case 'x':
                    if (lvf == 0)
                        {
                        p16 = *(int16_t *)varhead;
                        varhead = (int8_t *)varhead + sizeof(int16_t);
                        if (upf == 0)
                            {
                            n += outCStr(strlwr(itoh(p16)));
                            }
                        else
                            n += outCStr(itoh(p16));
                        }
                    else
                        {
                        p32 = *(int32_t *)varhead;
                        varhead = (int8_t *)varhead + sizeof(int32_t);
                        if (upf == 0)
                            {
                            n += outCStr(strlwr(ltoh(p32)));
                            }
                        else
                            n += outCStr(ltoh(p32));
                        }
                    break;
                default:
                    break;
                }
            p = i + 1;
            i++;
            }
        }

    return n;
}

void readBfr8(void * pv, void * bfr, size_t offset)
{
    memcpy(pv, ((int8_t *)bfr + offset), sizeof(int8_t));
}

void readBfr16(void * pv, void * bfr, size_t offset)
{
    memcpy(pv, ((int8_t *)bfr + offset), sizeof(int16_t));
}

void readBfr32(void * pv, void * bfr, size_t offset)
{
    memcpy(pv, ((int8_t *)bfr + offset), sizeof(int32_t));
}

DiskInfo * loadDrvInfo(int8_t DrvNum)
{
    static DiskInfo di;
    int16_t i16;
    int32_t i32;
    void * ptr;
    div_t lr;

    ptr = loadSec0(DrvNum);

    di.DrvNum = DrvNum;

    readBfr16(&di.SecInTrack, ptr, BPB_SEC_PER_TRACK_OFFSET);

    readBfr16(&di.HeadCnt, ptr, BPB_HEAD_CNT_OFFSET);

    readBfr16(&i16, ptr, BPB_TOTAL_SET16_OFFSET);

    if (i16 == 0)
        {
        readBfr32(&i32, ptr, BPB_TOTAL_SET32_OFFSET);
        di.SecCount = i32;
        }
    else
        {
        di.SecCount = i16;
        }

    readBfr32(&di.HiddenSecCnt, ptr, BPB_HIDDEN_SEC_OFFSET);

    readBfr8(&di.FatCount, ptr, BPB_FAT_COUNT_OFFSET);

    readBfr16(&di.FatSize, ptr, BPB_FAT_SIZE_OFFSET);

    readBfr16(&di.RootDirCnt, ptr, BPB_ROOT_DIR_CNT_OFFSET);

    readBfr16(&di.BytePerSec, ptr, BPB_BYTE_PER_SEC_OFFSET);

    readBfr8(&di.SecPerCluster, ptr, BPB_SEC_PER_CLUSTER_OFFSET);

    di.FatSecNum = di.HiddenSecCnt + 1;

    di.RootDirSecNum = di.FatSecNum + di.FatCount * di.FatSize;

    di.DataSecNum = di.RootDirSecNum + di.RootDirCnt / 16;

    if (di.RootDirCnt % 16 != 0)
        di.DataSecNum++;

    ldiv(di.SecCount, di.SecPerCluster, &lr);

    if (lr.quot < 0xFFF)
        di.Ftype = FAT12_TYPE, di.ClusterEnd = 0xFFF;
    else if (lr.quot < 0xFFFF)
        di.Ftype  = FAT16_TYPE, di.ClusterEnd = 0xFFFF;
    else
        di.Ftype  = FAT32_TYPE, di.ClusterEnd = 0xFFFFFFFF;

#if 0
    printf("disk info of DrvNum = %d\n", DrvNum);
    printf("FatCount %d\nFatSize %d\n", di.FatCount, di.FatSize);
    printf("disk ClusterSecNum %ld\ndisk RootDirSecNum %ld\ndisk DataSecNum %ld\n", di.ClusterSecNum, di.RootDirSecNum, di.DataSecNum);
#endif
    return &di;

}

void * loadSec0(int8_t DrvNum)
{
    static char buf[512];
    int16_t es_r;

    asm mov es_r, es;
    asm mov ax, ds;
    asm mov es, ax;
    asm mov ax, 0x0201;
    asm mov bx, offset buf;
    asm mov cx, 0x0001;
    asm mov dh, 0;
    asm mov dl, DrvNum;
    asm int 0x13;
    asm mov es, es_r;
    asm cmp ah, 0;
    asm jnz ldsFail;
    return buf;
ldsFail:
    return NULL;
}

void * loadSec(int16_t secNum, DiskInfo * dsi)
{
    static char buf[512];
    uint8_t sec, hd, trk, drvNum;
    int16_t es_r;

    trk = secNum / (dsi->SecInTrack * dsi->HeadCnt);
    sec = secNum % (dsi->SecInTrack * dsi->HeadCnt);
    hd = sec / dsi->SecInTrack;
    sec = (sec % dsi->SecInTrack) + 1;
    drvNum = dsi->DrvNum;

    asm mov es_r, es;
    asm mov ax, ds;
    asm mov es, ax;

    asm mov ax, 0x0201;
    asm mov bx, offset buf;
    asm mov cl, sec;
    asm mov ch, trk;
    asm mov dh, hd;
    asm mov dl, drvNum;
    asm int 0x13;
    asm mov es, es_r;
    asm cmp ah, 0;
    asm jnz ldsFail;
    return buf;
ldsFail:
    return NULL;
}


FileInfo * parseFileInfo(int8_t * rawData)
{
    static FileInfo fi;

    if (rawData == NULL)
        return NULL;

    memcpy(fi.Name, rawData, 11);

    fi.AttFlag = rawData[11];

    fi.ModifyTime = *(int16_t *)(rawData + 0x16);

    fi.ModifyDate = *(int16_t *)(rawData + 0x18);

    fi.FirstCluster= *(int16_t *)(rawData + 0x1A);

    fi.Length = *(int32_t *)(rawData + 0x1C);

    return &fi;
}

void DispAllFile(DiskInfo * di)
{
    int8_t * pb;
    int8_t fn[10], ext[5], flagstr[8];
    int16_t i, j, k, l, m;
    FileInfo * fi;

    if (di == NULL)
        return;

    strcpy(flagstr, "rhsldan");

    k = 0, j = 0;

    while (k < di->RootDirCnt)
        {
        pb = loadSec(di->RootDirSecNum + j, di);

        if (pb == NULL)
            return;

        for (i = 0; i < di->BytePerSec && k < di->RootDirCnt; i += FAT_FDT_SIZE, k++)
            {
            memset(fn, ' ', sizeof(fn));
            memset(ext, ' ', sizeof(ext));

            fn[9] = '\0', ext[4] = '\0';


            fi = parseFileInfo(pb + i);

            if (fi == NULL)
                continue;

            memcpy(fn, fi->Name, 8);
            memcpy(ext, fi->Name + 8, 3);

            if (fn[0] != ' ' && fn[0] != 0 && fn[0] != (int8_t)0xE5)
                {
                *strchr(fn, ' ') = '\0';
                *strchr(ext, ' ') = '\0';

                if ((fi->AttFlag & 0xF) == 0xF)
                    continue;

                if (strlen(ext) > 0)
                    l = printf("%s.%s", fn, ext);
                else
                    l = printf("%s", fn);

                for (m = 1; m < 16 - l; m++)
                    {
                    putCh(' ');
                    }

                printf("Attr:");

                for (l = 0; l < 7; l++)
                    {
                    if ((fi->AttFlag  >> l) & 1)
                        printf("%c", flagstr[l]);
                    else
                        printf("-");
                    }
                printf("\t%ld\n", fi->Length);
                }
            }
        j++;
        }
}


FileInfo * searchFile(int8_t * filename, DiskInfo * di)
{
    int8_t * pb;
    int8_t fn[10], ext[5], lbuf[16];
    int16_t i, j, k;
    static FileInfo * fi, ret;

    if (di == NULL || filename == NULL)
        return NULL;

    k = 0, j = 0;

    while (k < di->RootDirCnt)
        {
        pb = loadSec(di->RootDirSecNum + j, di);

        if (pb == NULL)
            return NULL;

        for (i = 0; i < di->BytePerSec && k < di->RootDirCnt; i += FAT_FDT_SIZE, k++)
            {
            memset(fn, ' ', sizeof(fn));
            memset(ext, ' ', sizeof(ext));

            fn[9] = '\0', ext[4] = '\0';


            fi = parseFileInfo(pb + i);

            if (fi == NULL)
                continue;

            memcpy(fn, fi->Name, 8);
            memcpy(ext, fi->Name + 8, 3);

            if (fn[0] != ' ' && fn[0] != 0 && fn[0] != (int8_t)0xE5)
                {
                *strchr(fn, ' ') = '\0';
                *strchr(ext, ' ') = '\0';

                if ((fi->AttFlag & 0x0F) == 0x0F)
                    continue;

                if (strlen(ext) > 0)
                    {
                    strcpy(lbuf, fn);
                    strcat(lbuf, ".");
                    strcat(lbuf, ext);
                    }
                else
                    strcpy(lbuf, fn);

                if (strcmp(lbuf, filename) == 0)
                    {
                    memcpy(&ret, fi, sizeof(ret));
                    ret.di = di;

                    return &ret;
                    }
                }
            }
        j++;
        }

    return NULL;
}

static int32_t ClsNumToSecNum(int32_t clsNum, DiskInfo * di)
{
    int32_t sec, SecPCls;
    mul_t mr;

    if (di == NULL)
        return -1;

    SecPCls = di->SecPerCluster;

    clsNum -= 3;

    lmul(clsNum, SecPCls, &mr);

    sec = mr.low + di->DataSecNum + 1;

    return sec;
}

int16_t readFile(void * bfr, FileInfo * fi, int32_t offset, size_t size)
{
    int32_t clsNum, co, coInc, ct, BPerS, ccs;
    int16_t sec16, cl, cf, ret;
    mul_t mr;
    div_t lr;
    DiskInfo * di;
    int8_t * ps, * pt, *p;
    int16_t i, j;

    if (fi == NULL || fi->di == NULL || bfr == NULL)
        return -1;

    di = fi->di;

    if (size == 0)
        return 0;

    if (offset >= fi->Length)
        return -1;

    clsNum = fi->FirstCluster;

    BPerS = di->BytePerSec;

    lmul(BPerS, di->SecPerCluster, &mr);

    coInc = mr.low;

    co = 0, cf = 0, ret = 0;

    while (clsNum != di->ClusterEnd)
        {
        if (offset >= co)
            cf = 1;

        if (cf)
            {
            sec16 = ClsNumToSecNum(clsNum, di) & 0xFFFF;

            for (i = 0, ct = 0; i < di->SecPerCluster; i++, sec16++, ct += BPerS)
                {
                p = loadSec(sec16, di);
                ccs = offset - co - ct;
                if (ccs < BPerS)
                    {
                    if (ccs > 0)
                        {
                        cl = (size > BPerS - ccs) ? (BPerS - ccs & 0xFFFF) : size;
                        memcpy(bfr, p + ccs, cl);
                        size -= cl;
                        }
                    else
                        {
                        cl = (size > BPerS) ? (BPerS & 0xFFFF) : size;
                        memcpy(bfr, p, cl);
                        size -= cl;
                        }
                    bfr = (int8_t *)bfr + cl;
                    ret += cl;
                    }

                if ((int16_t)size <= 0)
                    return cl;
                }
            }

        co += coInc;

        switch(di->Ftype)
            {
            case FAT12_TYPE:
                ccs = clsNum;
                ccs = lushr(ccs, 1);
                lmul(ccs, 3, &mr);
                ccs = mr.low;
                break;
            case FAT16_TYPE:
                ccs = clsNum;
                ccs = lushl(ccs, 1);
                break;
            case FAT32_TYPE:
                ccs = clsNum;
                ccs = lushl(ccs, 2);
                break;
            }
        ldiv(ccs, di->BytePerSec, &lr);
        sec16 = lr.quot + di->FatSecNum;
        p = loadSec(sec16, di);

        ct = *(int32_t *)(p + (lr.rem & 0xFFFF));

        switch(di->Ftype)
            {
            case FAT12_TYPE:
                if (clsNum & 1)
                    {
                    clsNum = lushr(ct, 12) & 0xFFF;
                    }
                else
                    {
                    clsNum = ct & 0xFFF;
                    }
                break;
            case FAT16_TYPE:
                clsNum = ct & 0xFFFF;
                break;
            case FAT32_TYPE:
                clsNum = ct;
                break;
            }
        }

    return cl;
}

static void * realPtrToBase16(int32_t x)
{
    int16_t ld, hd;

    ld = x & 0xFFFF;
    hd = (x >> 16) & 0xFFFF;

    hd = hd << 12;

    asm mov dx, hd;
    asm mov ax, ld;
}

static void clearMem()
{
    int32_t off;

    for (off = EXPORT_BASE; off < EXPORT_LIMIT; off += 512)
        {
        memset(realPtrToBase16(off), 0, 512);
        }
}

int16_t loadPE(int8_t * peName, DiskInfo * di)
{
    int8_t * pb, * pr;
    int32_t off, realPtr, segDescrBase;
    int16_t i, c;
    FileInfo fi, *fp;

    if (peName == NULL || di == NULL)
        return -1;

    fp = searchFile(peName, di);

    if (fp == NULL)
        {
        printf("File not found %s\n", peName);
        return -1;
        }

    memcpy(&fi, fp , sizeof(fi));

    off = 0;

    readFile(&idh, &fi, off, sizeof(idh));

    if (idh.e_lfanew == 0 || idh.e_magic != IMAGE_DOS_SIGNATURE)
        {
        printf("It is not a PE file\n");
        return -2;
        }

    off = idh.e_lfanew + SIZE_OF_NT_SIGNATURE;

    readFile(&ifh, &fi, off, sizeof(ifh));

    if (ifh.SizeOfOptionalHeader != sizeof(ioh))
        {
        printf("It is not a PE file\n");
        return -2;
        }

    off += sizeof(ifh);

    readFile(&ioh, &fi, off, sizeof(ioh));

    off += sizeof(ioh);

    printf("File Alignment: %lx\n", ioh.FileAlignment);

    printf("Segment Alignment: %lx\n", ioh.SectionAlignment);

    printf("Segment information:\n");

    segDescrBase = off;

    clearMem();

    for (i = 0; i < ifh.NumberOfSections; i++)
        {
        off =  segDescrBase + sizeof(pish) * i;

        readFile(&pish, &fi, off, sizeof(pish));

        printf("Segment name: %s", pish.Name);

        if (strcmp(pish.Name, ".text") == 0
            || strcmp(pish.Name, ".data") == 0
            || strcmp(pish.Name, ".bss") == 0
            || strcmp(pish.Name, ".rdata") == 0)
            {
            printf("\nSegment base: %lx\n", pish.VirtualAddress + ioh.ImageBase);

            printf("Segment export to: %lx\n", pish.VirtualAddress + ioh.ImageBase);

            realPtr = 0;

            off = pish.PointerToRawData;

            while (pish.SizeOfRawData > 0)
                {
                pr = realPtrToBase16(pish.VirtualAddress + ioh.ImageBase + realPtr);

                readFile(pr, &fi, off, (size_t)ioh.FileAlignment);

                off += ioh.FileAlignment;

                realPtr += ioh.FileAlignment;

                pish.SizeOfRawData -= ioh.FileAlignment;
                }
            }
        else
            {
            printf("(Ignore Segment)\n");
            }
        }

    printf("Loading PE file <%s> successfully\n", peName);

    return 0;
}

